from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
# CENTERNET_PATH = '../CenterNet/src'
# sys.path.insert(0, CENTERNET_PATH)
import _init_paths
import os
import cv2
import numpy as np
import torch
from opts import opts
from detectors.detector_factory import detector_factory
from datasets.dataset_factory import get_dataset


def preprocess(file_path, bin_path):
    opt = opts().parse('{} --load_model {}'.format('ctdet', '../model/ctdet_coco_dla_2x.pth').split(' ')) 
    Dataset = get_dataset(opt.dataset, opt.task)
    opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
    os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
    Detector = detector_factory[opt.task]
    detector = Detector(opt)
    in_files = os.listdir(file_path)
    if not os.path.exists(bin_path):
        os.makedirs(bin_path)
    i = 0
    for file in sorted(in_files):   
        i = i + 1
        print(file, "===", i)
        image = cv2.imread(os.path.join(file_path, file))
        for scale in opt.test_scales:
            images, meta = detector.pre_process(image, scale, meta=None)
            img = np.array(images).astype(np.float32)
            img.tofile(os.path.join(bin_path, file.split('.')[0] + '.bin'))

        
if __name__ == "__main__":
    file_path = os.path.abspath(sys.argv[1])
    bin_path = os.path.abspath(sys.argv[2])
    preprocess(file_path, bin_path)
